from typing import List, Dict, Any
import time
from pydantic import BaseModel
import pandas as pd
from calc_manager.plugins_inner.mat_csv_inner import MatCSVReader
from calc_manager.plugins_inner.op2_inner import OP2Plugin
from calc_manager.plugins_inner.bdf_inner import BDFPlugin
from calc_manager.plugins_inner.storage_plugin import StoragePlugin


class Parameter(BaseModel):
    bdf_model: str
    op2_fp: str

"""
class Result(BaseModel):
    danger_infos: Dict[str, List[int]]
    margin_output: List[Any]
    stress_output: List[Any]
    miss_mids: List[int]
    miss_eid_op2: List[int]
"""
class Result(BaseModel):
    version: str
    timestamp: str
    groups: List[Any]


def refine(param: Parameter, bdf_plugin: BDFPlugin, op2_plugin: OP2Plugin, storage_plugin: StoragePlugin) -> Result:
    storage_plugin.get_file("dc_storage_plugin/mat_csv/fcy.csv", 'fcy.csv')
    df = pd.read_csv('fcy.csv')

    def get_mat_from_csv(mid):
        # 示例材料的csv数据第一列是材料id,第二列是需要读取的值
        for index, row in df.iterrows():
            if row.iloc[0] == mid:
                return row.iloc[1]

    # 获取所有精细网格
    refine_elements = bdf_plugin.get_elements(param.bdf_model, element_types=['ctria3', 'cquad4'], refine=True)
    print(f'refine elements num: {len(refine_elements)}')
    rough_elements = bdf_plugin.get_elements(param.bdf_model, element_types=['ctria3', 'cquad4'], refine=False)
    print(f'rough elements num: {len(rough_elements)}')

    refine_eids = [refine_element.eid for refine_element in refine_elements]
    rough_eids = [rough_element.eid for rough_element in rough_elements]
    # 计算所有单元z2面的包络
    envelops = op2_plugin.get_shell_element_envelope(param.op2_fp, eids=refine_eids, result_types=['stress'], location='z2')
    print(f'envelops num: {len(envelops)}')
    eid_envelop_dict = {}
    for envelop in envelops:
        eid_envelop_dict[envelop.eid] = envelop
    # 记录所有mat csv里缺失的mid
    miss_mids = set()
    # 记录所有op2里缺失的eid
    miss_eid_op2 = set()

    danger_infos = {}
    eid_coffe = {}
    margin_outputs = {}
    stress_outputs = {}
    for refine_element in refine_elements:
        mid = refine_element.property.mids[0]
        # 根据mat_id获取材料许用应力数据
        material_ultimate_stress = get_mat_from_csv(mid)
        if material_ultimate_stress is None:
            miss_mids.add(mid)
            continue
        envelop = eid_envelop_dict.get(refine_element.eid)
        if envelop is None:
            miss_eid_op2.add(refine_element.eid)
            continue
        if max(envelop.max) >= material_ultimate_stress:
            print(f"危险单元 {refine_element.eid}")
            element_stress_max = max([envelop.max[5], material_ultimate_stress / 3])
            coffe = material_ultimate_stress / element_stress_max - 1
            # print(f'coffe: {coffe}')
            eid_coffe[refine_element.eid] = coffe
            g70_name = refine_element.component.g70_name
            if coffe < 0:
                cid = refine_element.cid
                if g70_name not in danger_infos:
                    danger_infos[g70_name] = [cid]
                elif cid not in danger_infos[g70_name]:
                    danger_infos[g70_name].append(cid)

            if g70_name not in margin_outputs.keys():
                margin_outputs[g70_name] = {
                    "gname": g70_name,
                    "eids": [refine_element.eid],
                    "coffe": [coffe],
                    "num": 1 if coffe < 0 else 0,
                }
            else:
                margin_outputs[g70_name]['eids'].append(refine_element.eid)
                margin_outputs[g70_name]['coffe'].append(coffe)
                if coffe < 0:
                    margin_outputs[g70_name]['num'] = margin_outputs[g70_name]['num'] + 1

            if g70_name not in stress_outputs.keys():
                stress_outputs[g70_name] = {
                    "gname": g70_name,
                    "eids": [refine_element.eid],
                    "stress": [element_stress_max]
                }
            else:
                stress_outputs[g70_name]['eids'].append(refine_element.eid)
                stress_outputs[g70_name]['stress'].append(element_stress_max)

    # print(f"danger_infos: {danger_infos}")
    # print(f'margin_outputs: {margin_outputs}')
    # print(f'stress_outputs: {stress_outputs}')
    # print(f'miss mids: {miss_mids}')
    # print(f'miss op2 eid: {miss_eid_op2}')

    # return Result(danger_infos=danger_infos,
    #               margin_output=list(margin_outputs.values()),
    #               stress_output=list(stress_outputs.values()),
    #               miss_mids=list(miss_mids),
    #               miss_eid_op2=list(miss_eid_op2))
    groups = [{
        "name": "全局图",
        "type": "diff",
        "value": {
            "rough": rough_eids,
            "refine": refine_eids
        }
    }]
    for margin_output in list(margin_outputs.values()):
        groups.append({
            "name": f'{margin_output["gname"]}危险单元裕度图',
            "type": "vec1",
            "eid": margin_output["eids"],
            "value": margin_output["coffe"],
            "label": [],
            "scalarBarTitle": "裕度图",
            "axis": {
                "position": "left-top",
            },
        })
    for stress_output in list(stress_outputs.values()):
        groups.append({
            "name": f'{stress_output["gname"]}危险单元应力云图',
            "type": "vec1",
            "eid": stress_output["eids"],
            "value": stress_output["stress"],
            "label": [],
            "scalarBarTitle": "应力图",
            "axis": {
                "position": "left-top",
            },
        })

    return Result(version="v1",
                  timestamp=f"{time.time()}",
                  groups=groups)